#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <spu_intrinsics.h>
#include <spu_mfcio.h>

#include "fish.h"
#include "flocking.h"

/* Flocking Parameters */
const float C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const float C1 = 2000.0; // weight of vector towards center of world
const float C1 = 0.0;
const float C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const float C3 = 0.0; // weight of vector towards center of gravity
const float C4 = 0.05; // weight of random vector
const float C5 = 0.1; // weight of vector pointing to cog of local flockmates
const float R_CROWD2 = 0.00390625; // square of crowding radius
const float R_SIGHT2 = 0.015625; // square of sighting radius
const float R_SIGHT = 0.125; // sighting radius
const float BIRD_MASS = 1.0; // bird mass
const float R = 1.0; // repulsion constant
const float M = 1.0; // constant that converts vectors to acceleration
const float START_CONST = 0.5;
const float SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const float SCALE_FACTOR = 1.0;
float v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay; // vectors
float cogx, cogy; // center of gravity coordinates
float local_cogx, local_cogy;
float tx, ty; // translation vectors for the moving grid

fish_t* fish[BUCKETS_Y][BUCKETS_X];
int fish_bucket_n[BUCKETS_Y][BUCKETS_X];
int n_fish;

void interact_fish_bucket(int, int);
void move_fish(float, int, int);
float compute_norm (int, int);
void bounce_fish (fish_t*);

int main(unsigned long long spe, unsigned long long argp, unsigned long long envp)
{
  // DMA in the spe_params_t struct
  spe_params_t spe_params __attribute__((aligned(16)));
  spe_auxinfo_t spe_auxinfo __attribute__((aligned(16)));
  float dt = 0.0;
  int tag, bucket_x, bucket_y, offset_x, offset_y, bucket_number, n, m, n0, m0;

  //printf("(Z, Z): Initiating Parameter Transfer\n\tmfc_ea2h(argp): %lu\n\tmfc_ea2l(argp): %lu\n\tmfc_ea2h(envp): %lu\n\tmfc_ea2l(envp): %lu\n", 
  //	 mfc_ea2h(argp), mfc_ea2l(argp), mfc_ea2h(envp), mfc_ea2l(envp));

  // issue the transfer
  tag = 1;
  spu_mfcdma64(&spe_params, mfc_ea2h(argp), mfc_ea2l(argp), sizeof(spe_params_t), tag, MFC_GET_CMD);
  spu_mfcdma64(&spe_auxinfo, mfc_ea2h(envp), mfc_ea2l(envp), sizeof(spe_auxinfo_t), tag, MFC_GET_CMD);

  // wait for all data with tag == 1
  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);



  bucket_number = spe_auxinfo.bucket_number;
  n = bucket_number / BUCKETS_X;
  m = bucket_number % BUCKETS_X;

  //printf("(%d, %d): Parameter Transfer Completed. Parameters:\n\tbucket#: %d\n\tfish address array:\n\t\t%lu\t%lu\t%lu\t%lu\n\t\t%lu\t%lu\t%lu\t%lu\n\tfish_bucket_n array:\n\t\t%d\t%d\t%d\t%d\n\t\t%d\t%d\t%d\t%d\n", 
/*
	 n, m, bucket_number, 
	 mfc_ea2l(spe_params.fish_ea[0][0]),
	 mfc_ea2l(spe_params.fish_ea[0][1]),
	 mfc_ea2l(spe_params.fish_ea[0][2]),
	 mfc_ea2l(spe_params.fish_ea[0][3]),
	 mfc_ea2l(spe_params.fish_ea[1][0]),
	 mfc_ea2l(spe_params.fish_ea[1][1]),
	 mfc_ea2l(spe_params.fish_ea[1][2]),
	 mfc_ea2l(spe_params.fish_ea[1][3]),
	 spe_params.fish_bucket_n[0][0],
	 spe_params.fish_bucket_n[0][1],
	 spe_params.fish_bucket_n[0][2],
	 spe_params.fish_bucket_n[0][3],
	 spe_params.fish_bucket_n[1][0],
	 spe_params.fish_bucket_n[1][1],
	 spe_params.fish_bucket_n[1][2],
	 spe_params.fish_bucket_n[1][3] );
*/

  unsigned int seed = spe_auxinfo.random_seed*((unsigned int)(-1));

  srand48(seed); 

  // recall that our x direction is horizontal
  for( offset_x = -1; offset_x < 2; offset_x++ ) {
    for( offset_y = 0; offset_y < 2; offset_y++ ) {
      n0 = (n+offset_y+BUCKETS_Y) % BUCKETS_Y;
      m0 = (m+offset_x+BUCKETS_X) % BUCKETS_X;

      fish_bucket_n[n0][m0] = spe_params.fish_bucket_n[n0][m0];

      if( fish_bucket_n[n0][m0] != 0 ) {
	fish[n0][m0] = (fish_t*)memalign(16, (fish_bucket_n[n0][m0])*sizeof(fish_t));

	//printf("(%d, %d): Transfering %d fish from bucket (%d, %d)\n", n, m, fish_bucket_n[n0][m0], n0, m0);
	//fflush(stdout);
	spu_mfcdma64(fish[n0][m0], mfc_ea2h((unsigned long long)spe_params.fish_ea[n0][m0]), mfc_ea2l((unsigned long long)spe_params.fish_ea[n0][m0]), 
		     fish_bucket_n[n0][m0]*sizeof(fish_t), tag, MFC_GET_CMD);
      }
    }
  }

  // wait on tag group 1
  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);

  //printf("(%d, %d): Fish Transfer Completed\n", n, m);

  // by this time we have the fish array filled up, and the fish_bucket_n 
  // array filled up... now we can run the interact_fish_bucket function from
  // our old code just fine.

  // calculate all the fish's ax and ay's
  //printf("(%d, %d): Calling interact_fish_bucket(%d, %d)\n", n, m, n, m);
  interact_fish_bucket( n, m );
  //printf("(%d, %d): interact_fish_bucket(%d, %d) has completed\n", n, m, n, m);

  // reduce max_norm for our fish
  float spe_max_norm = (float)compute_norm( n, m );

  /* Now get global max_norm */
  
  // here we're going to pull a crazy trick to put a floating point
  // number in the PPE's mailbox (it's expecting an unsigned int!)
  union {
    float f;
    int i;
  } u;
  
  u.f = spe_max_norm;
    
  
  spu_write_out_mbox(u.i); // tee hee! *girlish laughter*
  
  //printf("(%d, %d): wrote %f to outbound mailbox\n", n, m, u.f);
  
  u.i = spu_read_in_mbox(); // a blocking call

  //printf("(%d, %d): read %f from inbound mailbox\n", n, m, u.f);

  dt = (float)(u.f); // this should be the dt calculated by the PPE

  //printf("I, %d, have received dt: %f\n", n*BUCKETS + m, dt);

  move_fish(dt, n, m);


  // now we need to write the fish back to main memory
  spu_mfcdma64(fish[n][m], mfc_ea2h(spe_params.fish_ea[n][m]), mfc_ea2l(spe_params.fish_ea[n][m]), 
	       fish_bucket_n[n][m]*sizeof(fish_t), tag, MFC_PUT_CMD);
  

  spu_writech(MFC_WrTagMask, 1 << 1);
  spu_mfcstat(MFC_TAG_UPDATE_ALL);

  return 0;
}


void
move_fish(float dt, int n, int m)
{
  int i;

  for(i = 0; i < fish_bucket_n[n][m]; i++) {
    fish_t *__restrict__ f = &fish[n][m][i];
    
    f->x  += dt * f->vx * SCALE_FACTOR;
    f->y  += dt * f->vy * SCALE_FACTOR;
    f->vx += dt * f->ax;
    f->vy += dt * f->ay;
    
    bounce_fish(f);
  }
}


/* Compute the bucket maximum of all velocities / accelerations.  */
float
compute_norm (int n, int m)
{
  int p;
  float max_norm = 0.0;

  for(p = 0; p < fish_bucket_n[n][m]; p++){
    max_norm = f_max(max_norm, fabs(fish[n][m][p].vx));
    max_norm = f_max(max_norm, fabs(fish[n][m][p].vy));
    max_norm = f_max(max_norm, fabs(fish[n][m][p].ax));
    max_norm = f_max(max_norm, fabs(fish[n][m][p].ay));
  }
  
  return max_norm;
}


void
interact_fish_bucket(int n, int m)
{
  int i, j, ni, mi, n0, m0;
  fish_t *__restrict__ wfish;
  fish_t const *__restrict__ rfish;
  float r2, r16, r32, r4, r, a, angle, vtempx, vtempy, dx, dy;
  int in_sight, in_crowd;

  wfish = (fish_t *__restrict__)fish[n][m];

  for(i = 0; i < fish_bucket_n[n][m]; i++){

    v0x = v0y = v1x = v1y = v2x = v2y = v3x = v3y = v4x = v4y = v5x = v5y = vx = vy = ax = ay = 0.0;
    in_sight = in_crowd = 0;
    local_cogx = 0.0;
    local_cogy = 0.0;

    // calculate random unit vector
    v4x = drand48() - 0.5;
    v4y = drand48() - 0.5;
    r2 = f_max(v4x * v4x + v4y * v4y, min_r2);
    r = sqrt(r2);
    v4x /= r;
    v4y /= r;

    // interact with fish in this bucket and the 5 neighbor buckets
    for (ni = 0; ni < 2; ni++) {
      for (mi = -1; mi < 2; mi++) {

        n0 = (n+ni+BUCKETS_Y) % BUCKETS_Y;
        m0 = (m+mi+BUCKETS_X) % BUCKETS_X;
        rfish = (fish_t const *__restrict__)fish[n0][m0];

        for(j = 0; j < fish_bucket_n[n0][m0]; j++){
          dx = rfish[j].x - wfish[i].x;
          if ( fabs( dx ) > WALL_SEP - R_SIGHT ) {
	    if( dx < 0 )
	      dx = fabs( dx ) - WALL_SEP;
	    else
	      dx = WALL_SEP - fabs( dx );
          }
          dy = rfish[j].y - wfish[i].y;
          if ( fabs( dy ) > WALL_SEP - R_SIGHT ) {
	    if( dy < 0 )
	      dy = fabs( dy ) - WALL_SEP;
	    else
	      dy = WALL_SEP - fabs( dy );
          }

          if (dx != 0 || dy != 0) { 
            r2 = f_max(dx * dx + dy * dy, min_r2);

            // figure out if the boid is in this boid's range of vision
            angle = fabs( atan2(dy, dx) - atan2(wfish[i].vy, wfish[i].vx) );

            if( (angle <= (SIGHT_ANGLE/2.0)) || (angle >= (2*M_PI - SIGHT_ANGLE/2.0)) ) {
	      
              // calculate crowding vector
              if ( r2 < R_CROWD2 ) {
                v0x -= dx;
                v0y -= dy;
                in_crowd++;
              }

              // calculate flocking vector
              if( r2 < R_SIGHT2 ) {
                v2x += rfish[j].vx;
                v2y += rfish[j].vy;

                local_cogx += rfish[j].x;
                local_cogy += rfish[j].y;

                in_sight++;
              }
            }
          }
        }   
      }
    }

    if( in_sight != 0 ) {
      // calculate vector pointing in the direction of the average velocity of nearby flockmates
      v2x /= in_sight;
      v2y /= in_sight;

      v2x = v2x - wfish[i].vx;
      v2y = v2y - wfish[i].vy;

      local_cogx /= in_sight;
      local_cogy /= in_sight;
      
      // calculate vector towards cog of local flockmates
      v5x = local_cogx - wfish[i].x;
      v5y = local_cogy - wfish[i].y;
    }

    if( in_crowd != 0 ) {
      v0x /= in_crowd;
      v0y /= in_crowd;
      
      r2 = v0x * v0x + v0y * v0y;

      // let the repulsion decrease with r4
      v0x /= r2*r2;
      v0y /= r2*r2;
    }

    // now we have v0 to v5...
    vx = C0*v0x + C1*v1x + C2*v2x + C3*v3x + C4*v4x + C5*v5x;
    vy = C0*v0y + C1*v1y + C2*v2y + C3*v3y + C4*v4y + C5*v5y;   


    ax = M*vx;
    ay = M*vy;
    

    wfish[i].ax = ax;
    wfish[i].ay = ay;
  }
}

/* keep fish inside moving grid */
void
bounce_fish (fish_t* fish)
{
  while (fish->x < LEFT_WALL || fish->x > RIGHT_WALL) {
    if (fish->x < LEFT_WALL) {
      fish->x = fish->x + WALL_SEP;
    }
    if (fish->x > RIGHT_WALL) {
      fish->x = fish->x - WALL_SEP;
    }
  }
  while (fish->y < LEFT_WALL || fish->y > RIGHT_WALL) {
    if (fish->y < LEFT_WALL) {
      fish->y = fish->y + WALL_SEP;
    }
    if (fish->y > RIGHT_WALL) {
      fish->y = fish->y - WALL_SEP;
    }
  }
}

